package pool

import (
	"log"
	"os"
	"sync"
	"time"
)

type chanExecutorPool struct {
	size int

	taskChan  chan func()
	sizeChan  chan int
	closeChan chan bool
	wg        sync.WaitGroup

	logger *log.Logger
}

func (ep *chanExecutorPool) worker() {
	defer func() {
		ep.wg.Done()
		ep.logger.Println("worker close", time.Now())
	}()

	ep.logger.Println("worker start", time.Now())

	for {
		select {
		case t := <-ep.taskChan:
			ep.logger.Println("doing work...", time.Now())
			func() {
				defer func() {
					if er := recover(); er != nil {
						ep.logger.Printf("task error: %v\n", er)
					}
				}()
				t()
			}()
		case c := <-ep.closeChan:
			if c {
				ep.logger.Println("worker closed")
				return
			}
		}
	}
}

func (ep *chanExecutorPool) changeSize() {
	defer ep.wg.Done()
	for v := range ep.sizeChan {
		ep.size += v
	}
}

func (ep *chanExecutorPool) increase(n int) {
	for i := 0; i < n; i++ {
		ep.wg.Add(1)
		go ep.worker()
		ep.sizeChan <- 1
	}
}

func (ep *chanExecutorPool) reduce(n int) {
	for i := 0; i < n; i++ {
		ep.closeChan <- true
		ep.sizeChan <- -1
	}
}

func (ep *chanExecutorPool) start() {
	ep.wg.Add(1)
	go ep.changeSize()
	ep.increase(ep.size)
}

func (ep *chanExecutorPool) Submit(t func()) {
	ep.taskChan <- t
}

func (ep *chanExecutorPool) Shutdown() {
	ep.reduce(ep.size)
	close(ep.closeChan)
	close(ep.sizeChan)
}

func (ep *chanExecutorPool) WaitAll() {
	ep.wg.Wait()
}

type ExecutorPool interface {
	Submit(t func())
	Shutdown()
	WaitAll()
}

func NewChanExecutorPool(size, taskCap int) ExecutorPool {
	executor := &chanExecutorPool{
		size:      size,
		sizeChan:  make(chan int),
		taskChan:  make(chan func(), taskCap),
		closeChan: make(chan bool, size+1),
		logger:    log.New(os.Stdout, "executor-", log.LUTC|log.Lshortfile),
	}

	executor.start()

	return executor
}
